# -*- coding:utf-8 -*-#
# @Time:2023/6/29 16:10
# @Author:Adong
# @Software:PyCharm




'''
wav转Mel demo
'''



import librosa
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from PIL import Image

def Mel_filter(y,sr):
    pic_size = 224
    n_fft = (pic_size-1)*2
    hop_length = pic_size-1
    mel_spect = librosa.feature.melspectrogram(y=y, sr=sr, n_fft=n_fft, hop_length=hop_length)
    noise_stft = librosa.stft(y, n_fft=n_fft, hop_length=hop_length, window="hamm")  # stft得到时频谱

    # fig = plt.figure()
    # ax1 = plt.subplot(121)
    # ax2 = plt.subplot(122)
    # librosa.display.specshow(np.log(abs(mel_spect)), y_axis='mel',  x_axis='time',ax=ax1)
    # librosa.display.specshow(np.log(abs(noise_stft)), y_axis='hz', x_axis='time',ax=ax2)
    # plt.show(block=True)

    normalize_tool = MinMaxScaler(feature_range=(0, 255))  # 导入最大最小归一化工具，设置归一化范围为(0,255)
    mel_spect = np.array(zero_padding(mel_spect))
    normalized_pic = normalize_tool.fit_transform(mel_spect)  # 执行归一化
    img = Image.fromarray(normalized_pic)  # 将array转为image
    img = img.convert('L')  # 将image转为灰度图
    # img = img.resize((256,256))
    img.save('zlpc.png')  # 保存灰度图
    return 0


def zero_padding(mel_spect):
    '''把128*224的Mel图居中零填充到224*224的尺寸'''
    mel_spect = mel_spect.tolist()
    zero = np.zeros(len(mel_spect[0])).tolist()
    while len(mel_spect) < 224:
        mel_spect.append(zero)
    return np.array(mel_spect)


if __name__ == '__main__':
    filepath = ['./data/wav_data_V5/normal/环境声响_正常运行_normal_01.wav',r'./data/wav_data_V5/jbfd/局部放电_byq_jbfd_01.wav',
                r'./data/wav_data_V5/zgz/重过载_byq_zgz_01.wav',r'./data/wav_data_V5/zjsd/组件松动_byq_jjsd_01.wav',
                r'./data/wav_data_V5/zlpc/直流偏磁_byq_zlpc_01.wav']
    # data / wav_data_V4 / jbfdzgz / 局部放电_byq_jbfd_重过载_byq_zgz_4.wav
    y,sr = librosa.load(filepath[4])
    Mel_filter(y[:49729],sr)



